package main

import (
	"encoding/binary"
	"fmt"
	"io"
)

// QUIC 隧道协议常量
const (
	// Magic 魔数标识 "QUIC"
	ProtocolMagic uint32 = 0x51554943

	// 协议版本 - 🚀 升级到v2支持连接池
	ProtocolVersion uint32 = 0x00000002

	// Header 大小（固定） - 🚀 扩展到144字节
	HeaderSize = 144

	// 消息类型
	MessageTypeRequest  uint32 = 1
	MessageTypeResponse uint32 = 2
	MessageTypeError    uint32 = 3
)

// TunnelMessage 隧道消息
type TunnelMessage struct {
	Magic        uint32       // 魔数
	Version      uint32       // 版本
	MessageType  uint32       // 消息类型
	ProtocolType ProtocolType // 🚀 新增：协议类型（HTTP/TCP/SSH等）
	TunnelID     string       // 隧道ID (36字节 UUID)
	RequestID    string       // 请求ID (36字节 UUID)
	ConnectionID string       // 🚀 新增：连接ID (36字节 UUID)
	BodyLength   uint64       // Body 长度
	Reserved     [12]byte     // 保留字段（12字节，未来扩展）
	Body         []byte       // 消息体
}

// EncodeHeader 编码消息头部（144字节固定）
func (m *TunnelMessage) EncodeHeader() []byte {
	header := make([]byte, HeaderSize)

	// Magic (4 bytes) [0-4]
	binary.BigEndian.PutUint32(header[0:4], m.Magic)

	// Version (4 bytes) [4-8]
	binary.BigEndian.PutUint32(header[4:8], m.Version)

	// MessageType (4 bytes) [8-12]
	binary.BigEndian.PutUint32(header[8:12], m.MessageType)

	// ProtocolType (1 byte) [12-13]
	header[12] = uint8(m.ProtocolType)

	// Reserved1 (3 bytes) [13-16] - 对齐，保留字段

	// TunnelID (36 bytes) [16-52]
	copy(header[16:52], []byte(m.TunnelID))

	// RequestID (36 bytes) [52-88]
	copy(header[52:88], []byte(m.RequestID))

	// ConnectionID (36 bytes) [88-124]
	copy(header[88:124], []byte(m.ConnectionID))

	// BodyLength (8 bytes) [124-132]
	binary.BigEndian.PutUint64(header[124:132], m.BodyLength)

	// Reserved (12 bytes) [132-144] - 保留字段
	copy(header[132:144], m.Reserved[:])

	return header
}

// DecodeHeader 解码消息头部
func DecodeHeader(header []byte) (*TunnelMessage, error) {
	if len(header) < HeaderSize {
		return nil, fmt.Errorf("header太短: %d bytes, 需要 %d bytes", len(header), HeaderSize)
	}

	msg := &TunnelMessage{}

	// Magic (4 bytes) [0-4]
	msg.Magic = binary.BigEndian.Uint32(header[0:4])
	if msg.Magic != ProtocolMagic {
		return nil, fmt.Errorf("无效的魔数: 0x%X, 期望: 0x%X", msg.Magic, ProtocolMagic)
	}

	// Version (4 bytes) [4-8]
	msg.Version = binary.BigEndian.Uint32(header[4:8])
	// 🚀 兼容v1和v2
	if msg.Version != ProtocolVersion && msg.Version != 0x00000001 {
		return nil, fmt.Errorf("不支持的协议版本: %d, 期望: %d", msg.Version, ProtocolVersion)
	}

	// MessageType (4 bytes) [8-12]
	msg.MessageType = binary.BigEndian.Uint32(header[8:12])

	// ProtocolType (1 byte) [12-13]
	msg.ProtocolType = ProtocolType(header[12])

	// Reserved1 (3 bytes) [13-16] - 忽略

	// TunnelID (36 bytes) [16-52]
	tunnelIDBytes := header[16:52]
	msg.TunnelID = string(tunnelIDBytes[:36])

	// RequestID (36 bytes) [52-88]
	requestIDBytes := header[52:88]
	msg.RequestID = string(requestIDBytes[:36])

	// ConnectionID (36 bytes) [88-124]
	connIDBytes := header[88:124]
	msg.ConnectionID = string(connIDBytes[:36])

	// BodyLength (8 bytes) [124-132]
	msg.BodyLength = binary.BigEndian.Uint64(header[124:132])

	// Reserved (12 bytes) [132-144]
	copy(msg.Reserved[:], header[132:144])

	return msg, nil
}

// Write 写入完整消息（Header + Body）
func (m *TunnelMessage) Write(w io.Writer) error {
	// 写入 Header
	header := m.EncodeHeader()
	if _, err := w.Write(header); err != nil {
		return fmt.Errorf("写入Header失败: %v", err)
	}

	// 写入 Body
	if m.Body != nil && len(m.Body) > 0 {
		// 🔥🔥🔥 关键优化：分块写入Body,避免大数据量导致Stream阻塞
		// 问题：一次性写入大量数据会导致QUIC流控窗口被填满,Write()阻塞等待对端读取
		// 解决：分块写入,每次64KB,减少单次写入阻塞时间
		const chunkSize = 64 * 1024 // 64KB

		for offset := 0; offset < len(m.Body); offset += chunkSize {
			end := offset + chunkSize
			if end > len(m.Body) {
				end = len(m.Body)
			}

			// 写入一个块
			if _, err := w.Write(m.Body[offset:end]); err != nil {
				return fmt.Errorf("写入Body失败(offset=%d): %v", offset, err)
			}
		}
	}

	return nil
}

// ReadMessage 从 Reader 读取完整消息
func ReadMessage(r io.Reader) (*TunnelMessage, error) {
	// 读取 Header
	headerBuf := make([]byte, HeaderSize)
	if _, err := io.ReadFull(r, headerBuf); err != nil {
		return nil, fmt.Errorf("读取Header失败: %v", err)
	}

	// 解码 Header
	msg, err := DecodeHeader(headerBuf)
	if err != nil {
		return nil, err
	}

	// 读取 Body
	if msg.BodyLength > 0 {
		// 限制最大Body大小（防止内存溢出）
		maxBodySize := uint64(100 * 1024 * 1024) // 100MB
		if msg.BodyLength > maxBodySize {
			return nil, fmt.Errorf("Body太大: %d bytes, 最大: %d bytes", msg.BodyLength, maxBodySize)
		}

		msg.Body = make([]byte, msg.BodyLength)
		if _, err := io.ReadFull(r, msg.Body); err != nil {
			return nil, fmt.Errorf("读取Body失败: %v", err)
		}
	}

	return msg, nil
}

// NewRequestMessage 创建请求消息
func NewRequestMessage(tunnelID, requestID string, body []byte) *TunnelMessage {
	return &TunnelMessage{
		Magic:       ProtocolMagic,
		Version:     ProtocolVersion,
		MessageType: MessageTypeRequest,
		TunnelID:    tunnelID,
		RequestID:   requestID,
		BodyLength:  uint64(len(body)),
		Body:        body,
	}
}

// NewResponseMessage 创建响应消息
func NewResponseMessage(tunnelID, requestID string, body []byte) *TunnelMessage {
	return &TunnelMessage{
		Magic:       ProtocolMagic,
		Version:     ProtocolVersion,
		MessageType: MessageTypeResponse,
		TunnelID:    tunnelID,
		RequestID:   requestID,
		BodyLength:  uint64(len(body)),
		Body:        body,
	}
}

// NewErrorMessage 创建错误消息
func NewErrorMessage(tunnelID, requestID string, errorCode uint64, errorMsg string) *TunnelMessage {
	var reserved [12]byte
	// 将errorCode写入Reserved的前8字节
	for i := 0; i < 8 && i < 12; i++ {
		reserved[i] = byte(errorCode >> (uint(i) * 8))
	}

	return &TunnelMessage{
		Magic:       ProtocolMagic,
		Version:     ProtocolVersion,
		MessageType: MessageTypeError,
		TunnelID:    tunnelID,
		RequestID:   requestID,
		BodyLength:  uint64(len(errorMsg)),
		Reserved:    reserved,
		Body:        []byte(errorMsg),
	}
}
